from PyQt6.QtCore import QAbstractTableModel, QModelIndex, Qt, QVariant
import time
import Global

class CommTableModel(QAbstractTableModel):
    def __init__(self, data=None, headers=None, parent=None):
        super().__init__(parent)
        self._data = data or []
        self._headers = headers or []
        self._checkable_columns = set()  # 可配置复选框的列集合
        self._checked_states = {}  # 存储复选框状态的字典

    def rowCount(self, parent=QModelIndex()):
        return len(self._data)

    def columnCount(self, parent=QModelIndex()):
        return len(self._headers) if self._headers else (len(self._data[0]) if self._data else 0)

    def data(self, index, role=Qt.ItemDataRole.DisplayRole):
        if not index.isValid() or not (0 <= index.row() < self.rowCount()):
            return QVariant()

        if role == Qt.ItemDataRole.DisplayRole or role == Qt.ItemDataRole.EditRole:
            # 处理列数不匹配的情况
            if index.column() >= len(self._data[index.row()]):
                return QVariant()
            # 直接返回原始数据，避免不必要的字符串转换
            return self._data[index.row()][index.column()]
        elif role == Qt.ItemDataRole.CheckStateRole and index.column() in self._checkable_columns:
            key = (index.row(), index.column())
            return Qt.CheckState.Checked if self._checked_states.get(key, False) else Qt.CheckState.Unchecked
        return QVariant()

    def headerData(self, section, orientation, role=Qt.ItemDataRole.DisplayRole):
        if role != Qt.ItemDataRole.DisplayRole:
            return QVariant()

        if orientation == Qt.Orientation.Horizontal and section < len(self._headers):
            return self._headers[section]
        elif orientation == Qt.Orientation.Vertical:
            return str(section + 1)

        return QVariant()


    def setCheckableColumns(self, columns):
        """设置哪些列显示复选框"""
        self._checkable_columns = set(columns)
        
    def flags(self, index):
        if not index.isValid():
            return Qt.ItemFlag.NoItemFlags
        
        flags = Qt.ItemFlag.ItemIsEnabled | Qt.ItemFlag.ItemIsSelectable | Qt.ItemFlag.ItemIsEditable
        if index.column() in self._checkable_columns:
            flags |= Qt.ItemFlag.ItemIsUserCheckable
        return flags
        
    def setChecked(self, index, checked):
        """自定义设置复选框状态函数"""
        if not index.isValid():
            return False
            
        # 只更新复选框状态，不影响单元格内容
        self.dataChanged.emit(index, index, [Qt.ItemDataRole.CheckStateRole])
        return True
        
    def isChecked(self, index):
        """自定义获取复选框状态函数"""
        if not index.isValid():
            return False
        key = (index.row(), index.column())
        return self._checked_states.get(key, False)
    def getCheckedRows(self):
        true_keys = [key for key, val in self._checked_states.items() if val is True]
        return true_keys

    def setData(self, index, value, role=Qt.ItemDataRole.DisplayRole):
        if not index.isValid():
            return False
            
        if role == Qt.ItemDataRole.DisplayRole:
            # 直接存储原始数据，避免不必要的类型转换
            self._data[index.row()][index.column()] = value
            self.dataChanged.emit(index, index, [role])
            return True
        elif role == Qt.ItemDataRole.CheckStateRole and index.column() in self._checkable_columns:
            key = (index.row(), index.column())
            self._checked_states[key] = value == Qt.CheckState.Checked.value
            self.dataChanged.emit(index, index, [role])
            return True
            
        return False

    def insertRows(self, row: int, count: int, parent: QModelIndex = QModelIndex()) -> bool:
        self.beginInsertRows(parent, row, row + count - 1)
        self._data[row:row] = [[""] * self.columnCount() for _ in range(count)]
        self.endInsertRows()
        return True
        
    def insertRowsWithData(self, row: int, data: list,parent: QModelIndex = QModelIndex()) -> bool:
        """批量插入带有数据的行
        Args:
            row: 插入行位置
            data: 要插入的数据列表
            parent: 父索引
            number_column: 自动编号列位置，-1表示不自动编号
        """
        count = len(data)
        self.beginInsertRows(parent, row, row + count - 1)
        self._data[row:row] = data
        self.endInsertRows()
        self._checked_states={}
        return True

    def removeRows(self, row, count, parent=QModelIndex()):
        self.beginRemoveRows(parent, row, row + count - 1)
        for _ in range(count):
            self._data.pop(row)
        self.endRemoveRows()
        return True

    def setHeaders(self, headers):
        self._headers = headers
        self.headerDataChanged.emit(Qt.Orientation.Horizontal, 0, len(headers) - 1)

    def setDataList(self, data):
        self.beginResetModel()
        self._data = data
        self.endResetModel()

    def clear(self):
        self.beginResetModel()
        self._data = []
        self.endResetModel()
        
    def appendRow(self, row_data):
        """添加一行数据"""
        self.beginInsertRows(QModelIndex(), self.rowCount(), self.rowCount())
        self._data.append(row_data)
        self.endInsertRows()
        
    def removeRow(self, row):
        """删除指定行数据"""
        if 0 <= row < self.rowCount():
            self.beginRemoveRows(QModelIndex(), row, row)
            self._data.pop(row)
            self.endRemoveRows()
            return True
        return False
    
    def batchSetData(self, indices_values, role=Qt.ItemDataRole.DisplayRole):
        """批量设置数据
        Args:
            indices_values: 列表，每个元素为(index, value)元组
            role: 数据角色
        """
        if not indices_values:
            return True
            
        # 收集需要更新的范围
        min_row = min(idx.row() for idx, _ in indices_values)
        max_row = max(idx.row() for idx, _ in indices_values)
        min_col = min(idx.column() for idx, _ in indices_values)
        max_col = max(idx.column() for idx, _ in indices_values)
        
        # 批量更新数据
        for index, value in indices_values:
            if not index.isValid():
                continue
                
            if role == Qt.ItemDataRole.DisplayRole:
                if index.row() < len(self._data) and index.column() < len(self._data[index.row()]):
                    self._data[index.row()][index.column()] = value
            elif role == Qt.ItemDataRole.CheckStateRole and index.column() in self._checkable_columns:
                key = (index.row(), index.column())
                self._checked_states[key] = value == Qt.CheckState.Checked.value
        
        # 发出数据变化信号
        top_left = self.index(min_row, min_col)
        bottom_right = self.index(max_row, max_col)
        self.dataChanged.emit(top_left, bottom_right, [role])
        return True
    
    def batchRemoveRows(self, rows):
        """批量删除行
        Args:
            rows: 要删除的行号列表，必须已按从大到小排序
        """
        if not rows:
            return True
            
        # 确保行号按从大到小排序
        sorted_rows = sorted(set(rows), reverse=True)
        
        # 批量删除操作
        if sorted_rows:
            min_row = min(sorted_rows)
            max_row = max(sorted_rows)
            self.beginRemoveRows(QModelIndex(), min_row, max_row)
            
            # 按从大到小的顺序删除行，避免索引变化影响
            for row in sorted_rows:
                if 0 <= row < len(self._data):
                    self._data.pop(row)
            
            self.endRemoveRows()
        
        # 直接清除所有复选框状态
        self._checked_states = {}
            
        return True